from typing import Optional, Tuple, List
from typeguard import check_argument_types
from pathlib import Path
import numpy as np
import librosa
import onnxruntime # 1.12.1
import logging
import os


class Encoder:
    def __init__(self, model_dir: str) -> None:
        self.model_dir = model_dir
        model_path = os.path.join(self.model_dir, "full", "conformer_encoder.onnx")
        # load weight
        providers = ["CUDAExecutionProvider"]#['CPUExecutionProvider']
        self.encoder = onnxruntime.InferenceSession(model_path, providers=providers)
    
    def __call__(self, speech, lengths) -> List:
        """Inference
        Args:
            data: Input speech data
        Returns:
            text, token, token_int, hyp
        """
        assert check_argument_types()

        # # check dtype
        # if speech.dtype != np.float32:
        #     speech = speech.astype(np.float32)

        # # data: (Nsamples,) -> (1, Nsamples)
        # speech = speech[np.newaxis, :]
        # # lengths: (1,)
        # lengths = np.array([speech.shape[1]]).astype(np.int64)
        # logging.debug(" ###### before Forward_Encoder: ", speech.shape, lengths.shape)

        # b. Forward Encoder
        encoder_out, encoder_out_lens = self.encoder_main(speech=speech, speech_length=lengths)
        if isinstance(encoder_out, tuple):
            encoder_out = encoder_out[0]
        assert len(encoder_out) == 1, len(encoder_out)

        return encoder_out, encoder_out_lens


    def encoder_main(self, speech: np.ndarray, speech_length: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:

        """Frontend + Encoder. Note that this method is used by asr_inference.py
        Args:
            speech: (Batch, Length, ...)
            speech_lengths: (Batch, )
        """

        encoder_out, encoder_out_lens = self.encoder.run(["encoder_out", "encoder_out_lens"], {
                    "feats": speech,
                    "feats_length": speech_length
                })

        # encoder_out = self.mask_output(encoder_out, encoder_out_lens)
        # if self.config.do_postencoder:
        #     encoder_out, encoder_out_lens = self.postencoder(
        #         encoder_out, encoder_out_lens
        #     )
        logging.debug(encoder_out, encoder_out_lens)
        return encoder_out, encoder_out_lens
